import math, json
from scipy.optimize import nnls
import numpy as n
import scipy as s
from pyflix.datasets import RatedDataset
from machine import *

F = 20
LAMBDA = .05
STORE = '/netflix/nnmf_mats/'

class NNMF_Solver(Machine):
  name = "nnmf"
  def __init__(self, tr, store, f=F):
    self.tr = tr
    self.store = store
    self.f = f
    self.iters = 0
    self.P = n.random.random([max(tr.userIDs()) + 1, f]) * 2 * (3.604/f)**.5
    self.Q = n.random.random([max(tr.movieIDs()) + 1, f]) * 2 * (3.604/f)**.5
  
  def iterate_user(self, quiet=False):
    print "iterating on users"
    n_user = 0
    f = self.f
    for uid in self.tr.userIDs():
      user = self.tr.user(uid)
      movies, ratings = zip(*list(user.iterValueRatings()))
      Q_u = self.Q[movies,:]
      a = n.vstack((Q_u, n.diag(n.array([math.sqrt(LAMBDA * len(ratings))]*f))))
      b = n.hstack((n.array(ratings), n.array([0]*f)))
      self.P[uid] = nnls(a, b)[0]
      n_user += 1
      if n_user % 10000 == 0 and not quiet:
        print " ", n_user, "of", max(self.tr.userIDs())
  
  def iterate_movie(self, quiet=False):
    print "iterating on movies"
    n_movie = 0
    f = self.f
    for mid in self.tr.movieIDs():
      movie = self.tr.movie(mid)
      users, ratings = zip(*list(movie.iterValueRatings()))
      P_m = self.P[users,:]
      a = n.vstack((P_m, n.diag(n.array([math.sqrt(LAMBDA * len(ratings))]*f))))
      b = n.hstack((n.array(ratings), n.array([0]*f)))
      self.Q[mid] = nnls(a, b)[0]
      n_movie += 1
      if n_movie % 1000 == 0 and not quiet:
        print " ", n_movie, "of", max(self.tr.movieIDs())
  
  def est_rating(self, mid, uid):
    return min(max(n.dot(self.Q[mid], self.P[uid]), 1.0), 5.0)
    
  def write_mats(self):
    n.savez(self.store, p=self.P, q=self.Q)
  
  def read_mats(self):
    data = n.load(self.store)
    self.P = data['p']
    self.Q = data['q']
  
  def do_iters(self, n):
    for i in range(n):
      print "begining iteration", (self.iters + 1)
      self.iterate_user(quiet=True)
      self.iterate_movie(quiet=True)
      self.calc_rmse(False)
      self.iters += 1
      